import torch
from torch.utils.data import DataLoader
from torch import nn
from datasets.dataset import TicketDataset
from models.price_predictor import PricePredictor
import json
from tqdm import tqdm

CONFIG_PATH = "movie-ticket-bidding/config/SqueezeExcitationModel.json"
CHECKPOINT_PATH = "movie-ticket-bidding/checkpoints/SqueezeExcitationModel/best_model.pt"
DATA_PATH = "movie-ticket-bidding/data/processed/test.pkl"
BATCH_SIZE = 32
DEVICE = torch.device("cuda")

# 加载数据
dataset = TicketDataset(DATA_PATH)
dataloader = DataLoader(dataset, batch_size=BATCH_SIZE)

# 加载模型结构
with open(CONFIG_PATH, "r") as f:
    model_cfg = json.load(f)["model"]

input_dim = dataset[0][0].shape[0]
model = PricePredictor(CONFIG_PATH, input_dim).to(DEVICE)
model.load_state_dict(torch.load(CHECKPOINT_PATH, weights_only=True))
model.eval()

# 评估
criterion = nn.MSELoss()
mae = nn.L1Loss()
total_mse, total_mae, total_correct = 0.0, 0.0, 0

with torch.no_grad():
    loop = tqdm(dataloader, desc="Testing")
    for x, y in loop:
        x, y = x.to(DEVICE), y.to(DEVICE)
        pred = model(x)
        total_mse += criterion(pred, y).item() * x.size(0)
        total_mae += mae(pred, y).item() * x.size(0)

        # 计算5%内准确率
        pred_vals = pred.squeeze().cpu().numpy()
        true_vals = y.squeeze().cpu().numpy()
        total_correct += sum(abs(p - t) / (t + 1e-8) <= 0.05 for p, t in zip(pred_vals, true_vals))

mse_score = total_mse / len(dataset)
mae_score = total_mae / len(dataset)
accuracy = total_correct / len(dataset) 

print(f"📊 测试集 MSE: {mse_score:.4f}, MAE: {mae_score:.4f}, Acc@5%: {accuracy * 100:.2f}%)")
